#!/usr/bin/env python3
"""
ArucoTracker 演示脚本

展示如何使用 ArucoTracker 进行ArUco标记跟踪
"""

import time
import sys
import os
import signal
from aruco_tracker import ArucoTracker, ArucoMarker


class ArucoTrackerDemo:
    """ArucoTracker 演示类"""
    
    def __init__(self):
        self.tracker = None
        self.running = True
        
        # 设置信号处理器
        signal.signal(signal.SIGINT, self._signal_handler)
    
    def _signal_handler(self, signum, frame):
        """信号处理器"""
        print("\n接收到中断信号，正在停止...")
        self.running = False
        if self.tracker:
            self.tracker.stop_tracking()
        sys.exit(0)
    
    def my_callback(self, tracker_instance):
        """
        用户自定义回调函数
        
        Args:
            tracker_instance: ArucoTracker实例
        """
        # 获取跟踪统计信息
        stats = tracker_instance.get_tracking_statistics()
        
        if stats['tracked_markers'] > 0:
            print(f"\r跟踪中: {stats['tracked_markers']}/{stats['total_markers']} 标记, "
                  f"FPS: {stats['fps']:.1f}, 检测: {stats['detection_count']}", end='')
            
            # 获取测试标记状态
            test_marker_status = tracker_instance.get_marker_status("test_marker")
            if test_marker_status and test_marker_status['is_tracked']:
                pos = test_marker_status['position']
                rot = test_marker_status['rotation']
                dist = test_marker_status['distance']
                
                print(f"\n测试标记状态:")
                print(f"  位置: ({pos[0]:.1f}, {pos[1]:.1f}, {pos[2]:.1f}) mm")
                print(f"  姿态: ({rot[0]:.1f}°, {rot[1]:.1f}°, {rot[2]:.1f}°)")
                print(f"  距离: {dist:.1f} mm")
                
                # 检查是否在运动
                marker = tracker_instance.registered_markers["test_marker"]
                if marker.is_moving():
                    print("  状态: 运动中")
                else:
                    print("  状态: 静止")
    
    def run_basic_demo(self, camera_index=0, calibration_file=None):
        """
        运行基本演示
        
        Args:
            camera_index (int): 相机索引
            calibration_file (str): 标定文件路径
        """
        print("="*60)
        print("ArucoTracker 基本演示")
        print("="*60)
        print("使用标记ID=0，大小=50mm进行测试")
        print("按 Ctrl+C 退出")
        print("="*60)
        
        try:
            # 创建跟踪器
            self.tracker = ArucoTracker(
                camera_index=camera_index,
                calibration_file=calibration_file,
                max_lost_frames=30
            )
            
            # 注册测试标记（使用默认的滑动窗口大小5）
            test_marker = ArucoMarker(
                name="test_marker",
                marker_id=0,
                size_mm=50.0
            )
            
            if not self.tracker.register_marker(test_marker):
                print("标记注册失败")
                return
            
            # 设置回调函数
            self.tracker.set_callback(self.my_callback)
            
            # 开始跟踪
            if not self.tracker.start_tracking():
                print("开始跟踪失败")
                return
            
            print("\n开始跟踪，请将ArUco标记(ID=0)放在相机前...")
            
            # 主循环
            try:
                while self.running and self.tracker.is_tracking:
                    time.sleep(0.1)
                    
            except KeyboardInterrupt:
                print("\n用户中断")
            
        except Exception as e:
            print(f"演示运行错误: {e}")
        
        finally:
            if self.tracker:
                self.tracker.stop_tracking()
            print("\n演示结束")
    
    def run_multi_marker_demo(self, camera_index=0, calibration_file=None):
        """
        运行多标记演示
        
        Args:
            camera_index (int): 相机索引
            calibration_file (str): 标定文件路径
        """
        print("="*60)
        print("ArucoTracker 多标记演示")
        print("="*60)
        print("使用多个ArUco标记进行测试")
        print("按 Ctrl+C 退出")
        print("="*60)
        
        try:
            # 创建跟踪器
            self.tracker = ArucoTracker(
                camera_index=camera_index,
                calibration_file=calibration_file,
                max_lost_frames=30
            )
            
            # 注册多个标记（使用默认的滑动窗口大小5）
            markers = [
                ArucoMarker("marker_0", 0, 50.0),
                ArucoMarker("marker_1", 1, 50.0),
                ArucoMarker("marker_2", 2, 50.0),
            ]
            
            for marker in markers:
                self.tracker.register_marker(marker)
            
            # 设置回调函数
            def multi_marker_callback(tracker_instance):
                stats = tracker_instance.get_tracking_statistics()
                print(f"\r跟踪: {stats['tracked_markers']}/{stats['total_markers']}, "
                      f"FPS: {stats['fps']:.1f}", end='')
                
                # 如果有多个标记被跟踪，计算相对位置
                tracked_markers = [
                    name for name, marker in tracker_instance.registered_markers.items()
                    if marker.is_tracked
                ]
                
                if len(tracked_markers) >= 2:
                    rel_pos = tracker_instance.get_relative_position(
                        tracked_markers[0], tracked_markers[1]
                    )
                    if rel_pos:
                        print(f"\n{tracked_markers[0]} -> {tracked_markers[1]}: "
                              f"距离 {rel_pos['relative_distance']:.1f}mm")
            
            self.tracker.set_callback(multi_marker_callback)
            
            # 开始跟踪
            if not self.tracker.start_tracking():
                print("开始跟踪失败")
                return
            
            print(f"\n开始跟踪 {len(markers)} 个标记...")
            
            # 主循环
            try:
                while self.running and self.tracker.is_tracking:
                    time.sleep(0.1)
                    
                    # 定期输出状态
                    stats = self.tracker.get_tracking_statistics()
                    if stats['tracked_markers'] > 0:
                        # 保存跟踪数据
                        self.tracker.save_tracking_data("tracking_data.json")
                    
            except KeyboardInterrupt:
                print("\n用户中断")
            
        except Exception as e:
            print(f"多标记演示运行错误: {e}")
        
        finally:
            if self.tracker:
                self.tracker.stop_tracking()
            print("\n多标记演示结束")


def main():
    """主函数"""
    print("ArucoTracker 演示程序")
    
    # 解析命令行参数
    camera_index = 0
    calibration_file = None
    demo_type = "basic"
    
    if len(sys.argv) > 1:
        if sys.argv[1] in ['basic', 'multi']:
            demo_type = sys.argv[1]
        else:
            try:
                camera_index = int(sys.argv[1])
            except ValueError:
                calibration_file = sys.argv[1]
    
    if len(sys.argv) > 2:
        if demo_type == sys.argv[1]:
            try:
                camera_index = int(sys.argv[2])
            except ValueError:
                calibration_file = sys.argv[2]
        else:
            calibration_file = sys.argv[2]
    
    if len(sys.argv) > 3:
        calibration_file = sys.argv[3]
    
    # 检查标定文件
    if calibration_file and not os.path.exists(calibration_file):
        print(f"警告: 标定文件不存在: {calibration_file}")
        calibration_file = None
    
    print(f"配置: 相机索引={camera_index}, 标定文件={calibration_file}, 演示类型={demo_type}")
    
    # 创建并运行演示
    demo = ArucoTrackerDemo()
    
    if demo_type == "multi":
        demo.run_multi_marker_demo(camera_index, calibration_file)
    else:
        demo.run_basic_demo(camera_index, calibration_file)


if __name__ == "__main__":
    main()
